{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "Step 1: Preparation \n",
        "\n",
        "Before you begin, make sure that you have set the runtime type to GPU  (Hardware acclerator: GPU)."
      ],
      "metadata": {
        "id": "Yt0hTC0eyJX-"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8DyiIfxxF3Ey"
      },
      "outputs": [],
      "source": [
        "ROOTPATH=\"/content/ailab/\" # root dir (a constant)\n",
        "ModelPath=ROOTPATH+\"Real-CUGAN/model/\" # model dir\n",
        "PendingPath=ROOTPATH+\"Real-CUGAN/pending/\" # input dir\n",
        "FinishPath=ROOTPATH+\"Real-CUGAN/finish/\" # output dir\n",
        "ModelName=\"up2x-latest-no-denoise.pth\" # default model\n",
        "Tile=4 #{0,1,2,3,4,auto}; the larger the number, the smaller the memory consumption\n",
        "\n",
        "# initialize environment\n",
        "!pip install torch opencv-python\n",
        "!git clone https://github.com/bilibili/ailab.git\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Step 2: Download model files from [here](https://drive.google.com/drive/folders/1UFgpV14uEAcgYvVw0fJuajzy1k7JIz6H) and save them - in a folder called `updated_weights` - under your Google Drive's root folder.\n",
        "\n",
        "Step 3: Put/upload your image(s) under `/content/aliab/Real-CUGAN/pending`.\n",
        "\n",
        "Then, run the following and choose the model you want."
      ],
      "metadata": {
        "id": "0zjBE_P0yQ5M"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tMWiga8GOhMz"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "os.mkdir(ModelPath)\n",
        "os.mkdir(PendingPath)\n",
        "os.mkdir(FinishPath)\n",
        "!cp -r /content/gdrive/MyDrive/updated_weights/* /content/ailab/Real-CUGAN/model/\n",
        "fileNames = os.listdir(PendingPath)\n",
        "print(\"Pending images:\")\n",
        "for i in fileNames:\n",
        "  print(\"\\t\"+i)\n",
        "fileNames = os.listdir(ModelPath)\n",
        "print(\"Model files available:\")\n",
        "for idx, i in enumerate(fileNames):\n",
        "  print(f\"{idx+1}. \\t {i}\")\n",
        "choice = int(input(\"Select model (leave blank for default)：\"))\n",
        "if choice:\n",
        "  ModelName=fileNames[choice-1]\n",
        "Amplification=ModelName[2] # amplifying ratio\n",
        "if (not os.path.isfile(ModelPath+ModelName)):\n",
        "  print(\"Warning: selected model file does not exist\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Step 4: Run the processing script."
      ],
      "metadata": {
        "id": "-_JcZ1RAyZtU"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v0ip7SPmI8nj"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "sys.path.append(\"/content/ailab/Real-CUGAN\")\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline\n",
        "import torch\n",
        "from torch import nn as nn\n",
        "from torch.nn import functional as F\n",
        "import os,sys,cv2\n",
        "import numpy as np\n",
        "from upcunet_v3 import RealWaifuUpScaler\n",
        "from time import time as ttime  \n",
        "fileNames = os.listdir(PendingPath)\n",
        "print(f\"using model {ModelPath+ModelName}\")\n",
        "upscaler = RealWaifuUpScaler(2, ModelPath+ModelName, half=True, device=\"cuda:0\")\n",
        "t0 = ttime()\n",
        "for i in fileNames:\n",
        "  torch.cuda.empty_cache()\n",
        "  try:\n",
        "    img = cv2.imread(PendingPath+i)[:, :, [2, 1, 0]]\n",
        "    result = upscaler(img,tile_mode=5,cache_mode=2,alpha=1)\n",
        "    cv2.imwrite(FinishPath+i,result[:, :, ::-1])\n",
        "  except RuntimeError as e:\n",
        "    print (i+\" FAILED\")\n",
        "    print (e)\n",
        "  else:\n",
        "    print(i+\" DONE\")\n",
        "    os.remove(PendingPath+i)\n",
        "t1 = ttime()\n",
        "print(\"time_spent\", t1 - t0)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Generated images should be under /content/aliab/Real-CUGAN/finish."
      ],
      "metadata": {
        "id": "7yYUAiIzzCaC"
      }
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "Real-CUGAN.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
